// sdf.c - SDF Driver Mock Implementation using GMSSL

#include "sdf.h"
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <gmssl/sm3.h>
#include <gmssl/sm2.h>
#include <gmssl/sm2_z256.h>
#include <gmssl/rand.h> 
#include <gmssl/error.h> 
// 移除了 #include <gmssl/err.h>，因为它在您的环境中不存在。

// 【关键修正 V7】：由于找不到声明 ERR_get_error() 的头文件，我们手动声明其原型，
// 以解决 C99 模式下的“隐式声明”错误。
unsigned long ERR_get_error(void);


// --- 内部状态和宏定义 ---
#define MOCK_HANDLE ((HANDLE)0xDEADBEEF)

// SDR 错误码 (只需保留 sdf.h 中未定义的特定错误码)
#define SDR_DECRYPT_FAIL 0x01000021 

#define SM3_DIGEST_LENGTH 32
#define SGD_SM3 0x00000004

// 模拟哈希上下文和状态
static SM3_CTX hash_context;
static int is_hashing = 0;
static int is_device_open = 0;

// ------------------------------------------------------------------
// 密钥存储和初始化
// ------------------------------------------------------------------
#define MAX_KEY_INDEX 4
#define MAX_SESSION_KEYS 16
#define MAX_SESSION_KEY_SIZE 32 

// 存储内部 ECC 密钥对
static SM2_KEY internal_sm2_keys[MAX_KEY_INDEX]; 
static int is_keys_initialized = 0;

// 存储会话密钥 (通过 extern 暴露给 test_sm2_key_mgmt.c)
uint8_t session_keys[16][MAX_SESSION_KEY_SIZE]; 
static int session_key_used[MAX_SESSION_KEYS] = {0};
static size_t session_key_lengths[MAX_SESSION_KEYS] = {0};

/**
 * 模拟内部密钥的生成和存储。
 */
static void initialize_mock_keys() {
    if (is_keys_initialized) return;

    if (sm2_key_generate(&internal_sm2_keys[1]) != 1) {
        fprintf(stderr, "Error: Mock SM2 Sign Key (Index 1) generation failed.\n");
    } else {
        printf("Mock SM2 Sign Key (Index 1) generated.\n");
    }
    if (sm2_key_generate(&internal_sm2_keys[2]) != 1) {
        fprintf(stderr, "Error: Mock SM2 Enc Key (Index 2) generation failed.\n");
    } else {
        printf("Mock SM2 Enc Key (Index 2) generated.\n");
    }
    
    is_keys_initialized = 1;
}

// ------------------------------------------------------------------
// 密钥句柄管理
// ------------------------------------------------------------------

/**
 * 分配一个新的会话密钥句柄。
 */
static HANDLE alloc_key_handle() {
    for (int i = 0; i < MAX_SESSION_KEYS; i++) {
        if (!session_key_used[i]) {
            session_key_used[i] = 1;
            // 句柄 = MOCK_HANDLE + index + 1
            return (HANDLE)((ULONG_PTR)MOCK_HANDLE + i + 1);
        }
    }
    return NULL;
}

/**
 * 从密钥句柄获取其内部数组索引。
 */
int get_key_index_from_handle(HANDLE hKey) {
    if (hKey == NULL) return -1;
    // 检查句柄是否是会话密钥句柄
    if ((ULONG_PTR)hKey <= (ULONG_PTR)MOCK_HANDLE) return -1;

    ULONG index = (ULONG)((ULONG_PTR)hKey - (ULONG_PTR)MOCK_HANDLE - 1);
    
    if (index >= MAX_SESSION_KEYS || !session_key_used[index]) {
        return -1;
    }
    return (int)index;
}


// ------------------------------------------------------------------
// 辅助函数：公钥点转换 (基于 SM2_KEY 结构体修正)
// ------------------------------------------------------------------

/**
 * 将 SDF ECCrefPublicKey 结构体转换为 GMSSL 期望的 65 字节未压缩公钥点数组 (0x04 || X || Y)。
 */
static LONG sdf_pubkey_to_sm2_point(const ECCrefPublicKey *sdf_pubkey, uint8_t pub_point[65]) {
    if (sdf_pubkey->bits != 256) return SDR_ALGNOTSUPPORT; 
    
    size_t coord_len = ECCref_MAX_LEN; // 32 bytes

    pub_point[0] = 0x04; // 未压缩点标识符
    memcpy(pub_point + 1, sdf_pubkey->x, coord_len);
    memcpy(pub_point + 1 + coord_len, sdf_pubkey->y, coord_len);
    
    return SDR_OK;
}

/**
 * 将 GMSSL SM2 密钥结构中的公钥点 (SM2_Z256_POINT) 转换为 SDF ECCrefPublicKey 结构。
 * 此时 SM2_Z256_POINT 被视为 64 字节的 X || Y 坐标块。
 */
static LONG gmssl_point_to_sdf_pubkey(const SM2_KEY *sm2_key, ECCrefPublicKey *sdf_pubkey) {
    if (sm2_key == NULL || sdf_pubkey == NULL) {
        return SDR_OUTARGERR;
    }
    
    sdf_pubkey->bits = 256;
    size_t coord_len = ECCref_MAX_LEN; // 32 bytes

    // 关键修正 V4：SM2_KEY.public_key 是 SM2_Z256_POINT 结构体，它包含 X, Y 坐标（共 64 字节）。
    // 将结构体的首地址视为 64 字节的 X || Y 数组。
    const BYTE *point_coords = (const BYTE *)&sm2_key->public_key;

    // GMSSL的SM2_Z256_POINT是X (32B) || Y (32B)
    memcpy(sdf_pubkey->x, point_coords, coord_len); // 提取 X 坐标
    memcpy(sdf_pubkey->y, point_coords + coord_len, coord_len); // 提取 Y 坐标
    
    return SDR_OK;
}


// ------------------------------------------------------------------
// --- 设备管理模拟 ---
// ------------------------------------------------------------------

LONG SDF_OpenDevice(HANDLE * phDeviceHandle) { 
    if (is_device_open) return SDR_COMMFAIL;
    *phDeviceHandle = MOCK_HANDLE; 
    is_device_open = 1; 

    // 在打开设备时初始化内部密钥
    initialize_mock_keys();

    return SDR_OK; 
}

LONG SDF_CloseDevice(HANDLE hDeviceHandle) { 
    if (hDeviceHandle != MOCK_HANDLE || !is_device_open) return SDR_COMMFAIL;
    is_device_open = 0; 
    return SDR_OK; 
}

LONG SDF_GetDeviceInfo(HANDLE hSessionHandle, DEVICEINFO *pstDeviceInfo) {
    if (hSessionHandle != MOCK_HANDLE || !is_device_open) return SDR_COMMFAIL;
    if (pstDeviceInfo == NULL) return SDR_OUTARGERR;

    memset(pstDeviceInfo, 0, sizeof(DEVICEINFO));
    strcpy(pstDeviceInfo->DeviceName, "GMSSL_SDF_MOCK");
    pstDeviceInfo->HashAlgAbility = SGD_SM3;
    return SDR_OK;
}

// ------------------------------------------------------------------
// --- SM3 接口封装 ---
// ------------------------------------------------------------------

LONG SDF_HashInit(HANDLE hSessionHandle, ULONG uiAlgID, ECCrefPublicKey *pucPublicKey, BYTE *pucID, ULONG uiIDLength) {
    if (hSessionHandle != MOCK_HANDLE || !is_device_open) return SDR_COMMFAIL;
    if (uiAlgID != SGD_SM3) return SDR_ALGNOTSUPPORT;
    if (is_hashing) return SDR_STEPERR;

    sm3_init(&hash_context);

    if (pucPublicKey != NULL && uiIDLength != 0) {
        uint8_t pub_point[65]; // SDF 公钥 (0x04 || X || Y)
        if (sdf_pubkey_to_sm2_point(pucPublicKey, pub_point) != SDR_OK) return SDR_INARGERR;

        BYTE Z_A[32]; 
        SM2_Z256_POINT temp_z_pub;

        // 关键修正 V4：只复制 64 字节的 X || Y 坐标到 SM2_Z256_POINT 结构体中
        memcpy((uint8_t *)&temp_z_pub, pub_point + 1, 64);
        
        // SM2 Z_A 计算：传入正确的 SM2_Z256_POINT 指针
        if (sm2_compute_z(Z_A, &temp_z_pub, (const char *)pucID, uiIDLength) != 1) {
             return SDR_ALGNOTSUPPORT; 
        }

        sm3_update(&hash_context, Z_A, 32);
    } else if (pucPublicKey != NULL || uiIDLength != 0) {
        return SDR_INARGERR;
    }

    is_hashing = 1;
    return SDR_OK;
}

LONG SDF_HashUpdate(HANDLE hSessionHandle, BYTE *pucData, ULONG uiDataLength) {
    if (hSessionHandle != MOCK_HANDLE || !is_device_open) return SDR_COMMFAIL;
    if (!is_hashing) return SDR_STEPERR;
    if (pucData == NULL && uiDataLength != 0) return SDR_INARGERR;
    
    if (uiDataLength > 0) {
        sm3_update(&hash_context, pucData, uiDataLength);
    }
    return SDR_OK;
}


LONG SDF_HashFinal(HANDLE hSessionHandle, BYTE *pucHash, ULONG *puiHashLength) {
    if (hSessionHandle != MOCK_HANDLE || !is_device_open) return SDR_COMMFAIL;
    if (!is_hashing) return SDR_STEPERR;
    if (pucHash == NULL || puiHashLength == NULL) return SDR_OUTARGERR;
    
    sm3_finish(&hash_context, pucHash); 
    *puiHashLength = SM3_DIGEST_LENGTH;
    is_hashing = 0;
    return SDR_OK;
}


// ------------------------------------------------------------------
// --- 销毁会话密钥 (SDF_DestroyKey) ---
// ------------------------------------------------------------------

LONG SDF_DestroyKey(HANDLE hSessionHandle, HANDLE hKeyHandle) {
    if (hSessionHandle != MOCK_HANDLE || !is_device_open) return SDR_COMMFAIL;
    int index = get_key_index_from_handle(hKeyHandle);
    
    if (index == -1) {
        printf("Error: Attempted to destroy invalid key handle 0x%lX.\n", (ULONG_PTR)hKeyHandle);
        return SDR_KEYNOTEXIST;
    }
    
    session_key_used[index] = 0;
    session_key_lengths[index] = 0;
    memset(session_keys[index], 0, MAX_SESSION_KEY_SIZE);
    
    printf("Key Handle 0x%lX (Index %d) destroyed successfully.\n", (ULONG_PTR)hKeyHandle, index);
    
    return SDR_OK;
}

// ------------------------------------------------------------------
// --- 导出ECC公钥 (SDF_ExportSignPublicKey_ECC/SDF_ExportEncPublicKey_ECC) ---
// ------------------------------------------------------------------

LONG SDF_ExportSignPublicKey_ECC(HANDLE hSessionHandle, ULONG uiKeyIndex, ECCrefPublicKey *pucPublicKey) {
    if (hSessionHandle != MOCK_HANDLE || !is_device_open) return SDR_COMMFAIL;
    if (pucPublicKey == NULL) return SDR_OUTARGERR;
    
    if (uiKeyIndex == 0 || uiKeyIndex >= MAX_KEY_INDEX) return SDR_KEYNOTEXIST; 
    
    // 检查私钥是否已初始化（判断是否存在 SM2 密钥）
    if (internal_sm2_keys[uiKeyIndex].private_key[0] == 0) {
        return SDR_KEYNOTEXIST;
    }
    
    return gmssl_point_to_sdf_pubkey(&internal_sm2_keys[uiKeyIndex], pucPublicKey);
}

LONG SDF_ExportEncPublicKey_ECC(HANDLE hSessionHandle, ULONG uiKeyIndex, ECCrefPublicKey *pucPublicKey) {
    return SDF_ExportSignPublicKey_ECC(hSessionHandle, uiKeyIndex, pucPublicKey);
}


// ------------------------------------------------------------------
// --- 生成会话密钥并用外部 ECC 公钥加密输出 (SDF_GenerateKeyWithEPK_ECC)
// ------------------------------------------------------------------

LONG SDF_GenerateKeyWithEPK_ECC(
    HANDLE hSessionHandle, 
    ULONG uiKeyBits, 
    ULONG uiAlgID, 
    ECCrefPublicKey *pucPublicKey, 
    ECCCipher *pucKey, 
    HANDLE *phKeyHandle) 
{
    if (hSessionHandle != MOCK_HANDLE || !is_device_open) return SDR_COMMFAIL;
    if (pucPublicKey == NULL || pucKey == NULL || phKeyHandle == NULL) return SDR_OUTARGERR;
    
    size_t key_len = uiKeyBits / 8;

    if (key_len == 0 || key_len > MAX_SESSION_KEY_SIZE) {
        return SDR_INARGERR; 
    }

    uint8_t session_key[MAX_SESSION_KEY_SIZE];
    uint8_t pub_point[65]; 
    size_t cipher_len;
    
    SM2_KEY temp_pub_key = {0}; // 初始化结构体
    
    // 1. 生成会话密钥 (K)
    if (rand_bytes(session_key, key_len) != 1) {
        return SDR_COMMFAIL; 
    }

    // 2. 转换外部公钥为 GMSSL 格式 (65字节点: 0x04 || X || Y)
    if (sdf_pubkey_to_sm2_point(pucPublicKey, pub_point) != SDR_OK) {
        return SDR_INARGERR;
    }

    // 3. 将公钥点填充到临时的 SM2_KEY 结构体中
    // 关键修正 V4：SM2_KEY.public_key 是 SM2_Z256_POINT 结构体 (64字节的 X || Y)
    // 从 pub_point 的第 2 字节开始复制 64 字节
    memcpy((BYTE *)&temp_pub_key.public_key, pub_point + 1, 64);

    // 4. 分配句柄并存储密钥 (暂时分配，加密失败时释放)
    HANDLE hKey = alloc_key_handle();
    if (hKey == NULL) {
        return SDR_COMMFAIL; 
    }
    int index = get_key_index_from_handle(hKey);

    // 5. 执行 SM2 加密封装 (C1 || C3 || C2 结构，GMSSL 内部处理)
    uint8_t sm2_cipher[1024]; 
    
    if (sm2_encrypt(
        &temp_pub_key, 
        session_key, key_len, 
        sm2_cipher, &cipher_len) != 1) 
    {
        // 修正：此时 ERR_get_error() 应该已经被手动声明
        fprintf(stderr, "SM2 Encrypt failed. GMSSL Error Code: 0x%lX\n", ERR_get_error());
        SDF_DestroyKey(hSessionHandle, hKey);
        return SDR_ALGNOTSUPPORT; 
    }

    // 6. 将 GMSSL 输出结构 (C1 || C3 || C2) 映射到 SDF ECCCipher 结构
    // C1: 65 bytes, C3: 32 bytes, C2: key_len bytes
    
    memcpy(pucKey->C1, sm2_cipher, 65); 
    memcpy(pucKey->C3, sm2_cipher + 65, 32); 
    
    pucKey->C2Len = (ULONG)key_len;
    memcpy(pucKey->C2, sm2_cipher + 65 + 32, key_len); 

    // 7. 存储会话密钥并返回句柄
    memcpy(session_keys[index], session_key, key_len);
    session_key_lengths[index] = key_len;
    *phKeyHandle = hKey;

    printf("Generated and Encapsulated Session Key (Len: %zu bytes). Handle: 0x%lX\n", key_len, (ULONG_PTR)hKey);
    return SDR_OK;
}

// ------------------------------------------------------------------
// --- 导入会话密钥并用内部 ECC 私钥解密 (SDF_ImportKeyWithISK_ECC)
// ------------------------------------------------------------------

LONG SDF_ImportKeyWithISK_ECC(
    HANDLE hSessionHandle, 
    ULONG uiISKIndex, 
    ECCCipher *pucKey, 
    HANDLE *phKeyHandle) 
{
    if (hSessionHandle != MOCK_HANDLE || !is_device_open) return SDR_COMMFAIL;
    if (pucKey == NULL || phKeyHandle == NULL) return SDR_OUTARGERR;
    if (uiISKIndex == 0 || uiISKIndex >= MAX_KEY_INDEX) return SDR_KEYNOTEXIST; 

    if (internal_sm2_keys[uiISKIndex].private_key[0] == 0) {
        return SDR_KEYNOTEXIST;
    }
    
    // 1. 从 SDF ECCCipher 结构体中重构 GMSSL 期望的密文结构 (C1 || C3 || C2)
    size_t key_len = pucKey->C2Len;
    if (key_len == 0 || key_len > MAX_SESSION_KEY_SIZE) {
        return SDR_INARGERR;
    }
    
    uint8_t sm2_cipher[1024];
    size_t cipher_len = 65 + 32 + key_len;
    
    memcpy(sm2_cipher, pucKey->C1, 65);
    memcpy(sm2_cipher + 65, pucKey->C3, 32);
    memcpy(sm2_cipher + 65 + 32, pucKey->C2, key_len); 

    uint8_t decrypted_key[MAX_SESSION_KEY_SIZE];
    size_t decrypted_len;
    
    // 2. 执行 SM2 解密
    if (sm2_decrypt(
        &internal_sm2_keys[uiISKIndex], 
        sm2_cipher, cipher_len, 
        decrypted_key, &decrypted_len) != 1) 
    {
        fprintf(stderr, "SM2 Decrypt failed. GMSSL Error Code: 0x%lX\n", ERR_get_error());
        return SDR_DECRYPT_FAIL; 
    }

    // 3. 校验解密长度
    if (decrypted_len != key_len) {
        fprintf(stderr, "Decrypted length mismatch. Expected %zu, got %zu.\n", key_len, decrypted_len);
        return SDR_DECRYPT_FAIL;
    }
    
    // 4. 分配句柄并存储密钥
    HANDLE hKey = alloc_key_handle();
    if (hKey == NULL) {
        return SDR_COMMFAIL; 
    }
    int index = get_key_index_from_handle(hKey);

    memcpy(session_keys[index], decrypted_key, decrypted_len);
    session_key_lengths[index] = decrypted_len;
    *phKeyHandle = hKey;

    printf("Decapsulated and Imported Session Key (Len: %zu bytes) using ISK Index %lu. Handle: 0x%lX\n", key_len, uiISKIndex, (ULONG_PTR)hKey);
    return SDR_OK;
}
